#pragma once

/*
 * asio thread pool
 */

#include <memory>
#include <vector>
#include <thread>
#include <assert.h>
#include "asio/asio.hpp"

namespace KCP {
	class CThreadPool : asio::noncopyable {
	public:
		explicit CThreadPool(size_t pool_size) : m_nextIndex(0) {
			auto defaultSize = pool_size;
			if (defaultSize <= 0) {
				defaultSize = std::thread::hardware_concurrency();
			}
			if (defaultSize <= 0) {
				defaultSize = 1;
			}
			for (size_t i = 0; i < defaultSize; ++i) {
				auto io_service = std::make_shared<asio::io_context>();
				auto work = std::make_shared<asio::io_context::work>(*io_service);
				m_vecIOContexts.emplace_back(io_service);
				m_vecWorks.emplace_back(work);
			}
		}

		virtual ~CThreadPool() {
		}
		using IoContextPtr = std::shared_ptr<asio::io_context>;
		using WorkPtr = std::shared_ptr<asio::io_context::work>;
		using ThreadPtr = std::shared_ptr<std::thread>;

	public:
		// Add task to rand io service to run.
		template<typename F, typename... Args>
		void AddTask(F&& f, Args&& ... args) {
			GetIoContext().post(std::bind(std::forward<F>(f), std::forward<Args>(args)...));
		}

		void Join() {
			for (size_t i = 0; i < m_vecIOContexts.size(); i++) {
				m_vecThreads[i]->join();
			}
		}

		void Stop() {
			for (size_t i = 0; i < m_vecIOContexts.size(); i++) {
				m_vecIOContexts[i]->stop();
			}
		}

		// start thread pool to run io service event.
		void Init() {
			for (size_t i = 0; i < m_vecIOContexts.size(); i++) {
				auto ios = m_vecIOContexts[i];
				auto thr = std::make_shared<std::thread>(
					// thread function, dealing ios run.
					[this, ios]() {
					   ios->run();
				    }
				);
				m_vecThreads.emplace_back(thr);
			}
		}

		// Get a io context object.
		asio::io_context& GetIoContext() {
			asio::io_context& ioContext = *m_vecIOContexts[m_nextIndex];
			if (++m_nextIndex == m_vecIOContexts.size()) {
				m_nextIndex = 0;
			}
			return ioContext;
		}

	protected:
		std::vector<IoContextPtr> m_vecIOContexts;
		std::vector<WorkPtr>      m_vecWorks;
		std::vector<ThreadPtr>    m_vecThreads;
		std::atomic_int           m_nextIndex;
	};
}